library(foreign)
library(AER)
source("joint-iv-utils.R")

N <- 11500
sims <- 1000
boots <- 1000


tau.hold <- matrix(NA, nrow = sims, ncol = 12)
se.hold <- matrix(NA, nrow = sims, ncol = 12)
rho.hold <- matrix(NA, nrow = sims, ncol = 9)
tsls.laje.hold <- rep(NA, sims)
tsls.laje.ci.hi <- rep(NA, sims)
tsls.laje.ci.lo <- rep(NA, sims)
tsls.lace.10.hold <- rep(NA, sims)
tsls.lace.10.ci.hi <- rep(NA, sims)
tsls.lace.10.ci.lo <- rep(NA, sims)
tsls.lace.20.hold <- rep(NA, sims)
tsls.lace.20.ci.hi <- rep(NA, sims)
tsls.lace.20.ci.lo <- rep(NA, sims)
tsls.lace.11.hold <- rep(NA, sims)
tsls.lace.11.ci.hi <- rep(NA, sims)
tsls.lace.11.ci.lo <- rep(NA, sims)
tsls.lace.21.hold <- rep(NA, sims)
tsls.lace.21.ci.hi <- rep(NA, sims)
tsls.lace.21.ci.lo <- rep(NA, sims)
tsls.laie.hold <- rep(NA, sims)
tsls.laie.ci.hi <- rep(NA, sims)
tsls.laie.ci.lo <- rep(NA, sims)

cprobs <- rep(0, times = 9)
cc.effects <- c(1, 0.2, 0.1, -0.1, 0.2, 0.2, -0.4, 0.2, 0.7)
names(cc.effects) <- c("cc", "ca", "cn", "ac", "aa", "an", "nc", "na", "nn")
names(cprobs) <- c("cc", "ca", "cn", "ac", "aa", "an", "nc", "na", "nn")

cprobs["cc"] <- 0.1
cprobs["cn"] <- 0.2
cprobs["nc"] <- 0.2
cprobs["nn"] <- 0.5


agg.lace1 <- cc.effects["cc"] * (cprobs["cc"]/sum(cprobs[c("cc", "cn")])) + cc.effects["cn"] * (cprobs["cn"]/sum(cprobs[c("cc", "cn")]))
agg.lace2 <- cc.effects["cc"] * (cprobs["cc"]/sum(cprobs[c("cc", "nc")])) + cc.effects["nc"] * (cprobs["nc"]/sum(cprobs[c("cc", "nc")]))

true.tau <- c(1, 2, 1, 2, 0.1, 0, -0.4, 0, 3, 1, agg.lace1, agg.lace2)



set.seed(12345)
for (s in 1:sims) {
  if ((s %% 10) == 0) cat("simulation ", s, " out of ", sims, "\n")
  ## categories:
  ## cc, ca, cn, aa, ac, an, nn, nc, na
  c1 <- t(rmultinom(N, 1, cprobs))
  colnames(c1) <- c("cc", "ca", "cn", "ac", "aa", "an", "nc", "na", "nn")

  z1 <- sample(rep(c(rep(0, times = 9),1), N/10))
  z2 <- sample(rep(c(rep(0, times = 4),1), N/5))

  d1 <- rowSums(c1[,c("cc", "ca", "cn")])*z1 + rowSums(c1[,c("aa", "ac", "an")])
  d2 <- rowSums(c1[,c("cc", "ac", "nc")])*z2 + rowSums(c1[,c("aa", "ca", "na")])

  beta.0 <- rnorm(N, rowSums(t(cc.effects * t(c1))), 1)
  beta.d1 <- rnorm(N, rowSums(t(cc.effects * t(c1))), 1)
  beta.d2 <- rnorm(N, rowSums(t(cc.effects * t(c1))), 1)
  beta.int <- rnorm(N, rowSums(t(cc.effects * t(c1))), 1)
  epsilon <- rnorm(N, 0, 1)

  y00 <- beta.0 + epsilon
  y10 <- beta.0 + beta.d1 + epsilon
  y01 <- beta.0 + beta.d2 + epsilon
  y11 <- beta.0 + beta.d1 + beta.d2 + beta.int + epsilon

  y <- y00 + d1 * (y10-y00) + d2 * (y01-y00) + d1 * d2 * ((y11 - y01) - (y10 - y00))

  out <- joint.iv(y, d1, d2, z1, z2)
  tau.hold[s,] <- out$tau
  se.hold[s,] <- out$se
  rho.hold[s,] <- out$rho

  tsls.mod <- ivreg(y ~ d1*d2|z1*z2)
  tmp <- vcov(tsls.mod)
  tsls.lace.10.hold[s] <- coef(tsls.mod)["d1"]
  this.se <- sqrt(tmp["d1", "d1"])
  tsls.lace.10.ci.hi[s] <- tsls.lace.10.hold[s] + qt(0.975, df = N-4) * this.se
  tsls.lace.10.ci.lo[s] <- tsls.lace.10.hold[s] - qt(0.975, df = N-4) * this.se

  tsls.lace.20.hold[s] <- coef(tsls.mod)["d2"]
  this.se <- sqrt(tmp["d2", "d2"])
  tsls.lace.20.ci.hi[s] <- tsls.lace.20.hold[s] + qt(0.975, df = N-4) * this.se
  tsls.lace.20.ci.lo[s] <- tsls.lace.20.hold[s] - qt(0.975, df = N-4) * this.se

  tsls.laie.hold[s] <- coef(tsls.mod)["d1:d2"]
  this.se <- sqrt(tmp["d1:d2", "d1:d2"])
  tsls.laie.ci.hi[s] <- tsls.laie.hold[s] + qt(0.975, df = N-4) * this.se
  tsls.laie.ci.lo[s] <- tsls.laie.hold[s] - qt(0.975, df = N-4) * this.se


  tsls.lace.11.hold[s] <- coef(tsls.mod)["d1"] + coef(tsls.mod)["d1:d2"]
  this.se <- sqrt(tmp["d1","d1"] + tmp["d1:d2","d1:d2"] + 2*tmp["d1","d1:d2"])
  tsls.lace.11.ci.hi[s] <- tsls.lace.11.hold[s] + qt(0.975, df = N-4) * this.se
  tsls.lace.11.ci.lo[s] <- tsls.lace.11.hold[s] - qt(0.975, df = N-4) * this.se

  tsls.lace.21.hold[s] <- coef(tsls.mod)["d2"] + coef(tsls.mod)["d1:d2"]
  this.se <- sqrt(tmp["d2","d2"] + tmp["d1:d2","d1:d2"] + 2*tmp["d2","d1:d2"])
  tsls.lace.21.ci.hi[s] <- tsls.lace.21.hold[s] + qt(0.975, df = N-4) * this.se
  tsls.lace.21.ci.lo[s] <- tsls.lace.21.hold[s] - qt(0.975, df = N-4) * this.se

  tsls.laje.hold[s] <- coef(tsls.mod)["d1"] + coef(tsls.mod)["d2"] + coef(tsls.mod)["d1:d2"]
  this.se <- tmp["d1","d1"] + tmp["d2","d2"] + tmp["d1:d2","d1:d2"] + 2*tmp["d1","d2"] + 2*tmp["d1","d1:d2"] + 2*tmp["d2","d1:d2"]
  tsls.laje.ci.hi[s] <- tsls.laje.hold[s] + qt(0.975, df = N-4) * this.se
  tsls.laje.ci.lo[s] <- tsls.laje.hold[s] - qt(0.975, df = N-4) * this.se
}

true.hold <- matrix(true.tau, nrow = sims, ncol = length(true.tau), byrow = TRUE)
tau.bias <- colMeans(tau.hold) - true.tau
tau.s.ses <- apply(tau.hold, 2, sd)
tau.se.bias <- colMeans(se.hold) - tau.s.ses
tau.rmse <- sqrt(colMeans((tau.hold - true.hold)^2))
tau.ci.hi <- tau.hold + 1.96 * se.hold
tau.ci.lo <- tau.hold - 1.96 * se.hold
tmp <- tau.ci.lo < true.hold & tau.ci.hi > true.hold
tau.cov <- colMeans(tau.ci.lo < true.hold & tau.ci.hi > true.hold)

names(tau.bias) <- names(tau.s.ses) <- names(tau.se.bias) <- names(tau.rmse) <- names(tau.ci.hi) <- names(tau.ci.lo) <- names(tau.cov) <- names(out$tau)

tau.bias["laie"]
tau.rmse["laie"]
tau.cov["laie"]

tau.bias["laje"]
tau.rmse["laje"]
tau.cov["laje"]

mean(tsls.laie.hold - 1)
sqrt(mean((tsls.laie.hold - 1)^2))
